#include "StdAfx.h"
#include "GDT.h"
#include "console.h"

// Each define here is for a specific flag in the descriptor.
// Refer to the intel documentation for a description of what each one does.
#define SEG_DESCTYPE(x)  ((x) << 0x04) // Descriptor type (0 for system, 1 for code/data)
#define SEG_PRES(x)      ((x) << 0x07) // Present
#define SEG_SAVL(x)      ((x) << 0x0C) // Available for system use
#define SEG_LONG(x)      ((x) << 0x0D) // Long mode
#define SEG_SIZE(x)      ((x) << 0x0E) // Size (0 for 16-bit, 1 for 32)
#define SEG_GRAN(x)      ((x) << 0x0F) // Granularity (0 for 1B - 1MB, 1 for 4KB - 4GB)
#define SEG_PRIV(x)     (((x) &  0x03) << 0x05)   // Set privilege level (0 - 3)
 
#define SEG_DATA_RD        0x00 // Read-Only
#define SEG_DATA_RDA       0x01 // Read-Only, accessed
#define SEG_DATA_RDWR      0x02 // Read/Write
#define SEG_DATA_RDWRA     0x03 // Read/Write, accessed
#define SEG_DATA_RDEXPD    0x04 // Read-Only, expand-down
#define SEG_DATA_RDEXPDA   0x05 // Read-Only, expand-down, accessed
#define SEG_DATA_RDWREXPD  0x06 // Read/Write, expand-down
#define SEG_DATA_RDWREXPDA 0x07 // Read/Write, expand-down, accessed
#define SEG_CODE_EX        0x08 // Execute-Only
#define SEG_CODE_EXA       0x09 // Execute-Only, accessed
#define SEG_CODE_EXRD      0x0A // Execute/Read
#define SEG_CODE_EXRDA     0x0B // Execute/Read, accessed
#define SEG_CODE_EXC       0x0C // Execute-Only, conforming
#define SEG_CODE_EXCA      0x0D // Execute-Only, conforming, accessed
#define SEG_CODE_EXRDC     0x0E // Execute/Read, conforming
#define SEG_CODE_EXRDCA    0x0F // Execute/Read, conforming, accessed
 
#define GDT_CODE_PL0 SEG_DESCTYPE(1) | SEG_PRES(1) | SEG_SAVL(0) | \
                     SEG_LONG(0)     | SEG_SIZE(1) | SEG_GRAN(1) | \
                     SEG_PRIV(0)     | SEG_CODE_EXRD
 
#define GDT_DATA_PL0 SEG_DESCTYPE(1) | SEG_PRES(1) | SEG_SAVL(0) | \
                     SEG_LONG(0)     | SEG_SIZE(1) | SEG_GRAN(1) | \
                     SEG_PRIV(0)     | SEG_DATA_RDWR
 
#define GDT_CODE_PL3 SEG_DESCTYPE(1) | SEG_PRES(1) | SEG_SAVL(0) | \
                     SEG_LONG(0)     | SEG_SIZE(1) | SEG_GRAN(1) | \
                     SEG_PRIV(3)     | SEG_CODE_EXRD
 
#define GDT_DATA_PL3 SEG_DESCTYPE(1) | SEG_PRES(1) | SEG_SAVL(0) | \
                     SEG_LONG(0)     | SEG_SIZE(1) | SEG_GRAN(1) | \
                     SEG_PRIV(3)     | SEG_DATA_RDWR

#define GDT_GET_LIMIT()		((m_lastEntryId + 1) * sizeof(QWORD))

GDT::GDT(void)
{
	m_isActivated = false;
	m_pIdt        = NULL;
	m_lastEntryId = 0;
	
	xmemset_d(m_entries, (DWORD)0, sizeof(m_entries));

	m_gdt.Base  = (DWORD)m_entries;
	m_gdt.Limit = GDT_GET_LIMIT(); 

	m_ds = 0;
	m_es = 0;
	m_gs = 0;
	m_fs = 0;

	// predefined values
	AddCodeEntry(0, 0xffffffff, 0);				// sel 0x08 is code
	AddDataEntry(0, 0xffffffff, false, 0);		// sel 0x10 is data
}

GDT::~GDT(void)
{
	if (m_isActivated)
	{
		// We can't unload activated GDT
		putline("Invalid Operation Exception: we can't unload GDT that is active");
		for(;;)
		{}
	}
}


bool IsValidSelector(QWORD* pSelectors, WORD selCount, WORD selToCheck, bool isCode = false)
{
	selToCheck >>= 3;
	if (selToCheck == 0 || selToCheck > selCount)
	{
		return false;
	}

	DWORD sel = pSelectors[selToCheck] >> 32;

	DWORD mask = (((SEG_DESCTYPE(1) | SEG_PRES(1) | SEG_PRIV(0)) <<  8) & 0x00F0FF00);

	if ((sel & mask) != mask)
	{
		return false;
	}

	return ((sel & 0x0f) > 0x07) && isCode;
}


void GDT::Activate()
{
	m_isActivated = true;

	DWORD gdtAddr = (DWORD)&m_gdt;

	__asm 
	{
		cli

		lgdt gdtAddr
	};

	if (IsValidSelector(m_entries, m_lastEntryId, m_ds))
	{
		WORD _ds = m_ds;
		__asm mov ds, _ds;
	}

	if (IsValidSelector(m_entries, m_lastEntryId, m_es))
	{
		WORD _es = m_es;
		__asm mov es, _es;
	}

	if (IsValidSelector(m_entries, m_lastEntryId, m_fs))
	{
		WORD _fs = m_fs;
		__asm mov fs, _fs;
	}

	if (IsValidSelector(m_entries, m_lastEntryId, m_gs))
	{
		WORD _gs = m_gs;
		__asm mov gs, _gs;
	}

	if (m_pIdt)
	{
		m_pIdt->Activate();
	}
}

QWORD __inline GDT::CreateEntry(DWORD base, DWORD limit, WORD flag)
{
	QWORD descriptor;
 
    // Create the high 32 bit segment
    descriptor  =  limit       & 0x000F0000;         // set limit bits 19:16
    descriptor |= (flag <<  8) & 0x00F0FF00;         // set type, p, dpl, s, g, d/b, l and avl fields
    descriptor |= (base >> 16) & 0x000000FF;         // set base bits 23:16
    descriptor |=  base        & 0xFF000000;         // set base bits 31:24
 
    // Shift by 32 to allow for low part of segment
    descriptor <<= 32;
 
    // Create the low 32 bit segment
    descriptor |= base  << 16;                       // set base bits 15:0
    descriptor |= limit  & 0x0000FFFF;               // set limit bits 15:0

	return descriptor;
}


GDT GDT::Attach()
{
	PGDT_OBJ pGdt = NULL;

	__asm
	{
		sgdt pGdt;
	};

	GDT gdt;

	int items = (pGdt->Limit + 1) / sizeof(QWORD);

	if (items > MAX_GDT_ENTRIES)
	{
		items = MAX_GDT_ENTRIES;
		putline("WARNING: existing GDT is has to many elements - it will be trunced");
	}

	for (int i = 0; i < items; i ++)
	{
		gdt.m_entries[i] = ((QWORD*)pGdt->Base)[i];
	}

	gdt.m_lastEntryId = items;

	WORD  _ds, _es, _fs, _gs;
	__asm
	{
		mov _ds, ds
		mov _es, es
		mov _fs, es
	}

	gdt.m_ds = _ds;
	gdt.m_es = _es;
	gdt.m_fs = _fs;
	gdt.m_gs = _gs;

	return gdt;
}

void GDT::AssumeDS(WORD sel)
{
	if (m_isActivated)
	{
		__asm
		{
			mov ax, sel
			mov ds, sel
		}
	}
	else
	{
		m_ds = sel;
	}
}

void GDT::AssumeES(WORD sel)
{
	if (m_isActivated)
	{
		__asm
		{
			mov ax, sel
			mov es, sel
		}
	}
	else
	{
		m_es = sel;
	}
}


void GDT::AssumeFS(WORD sel)
{
	if (m_isActivated)
	{
		__asm
		{
			mov ax, sel
			mov fs, sel
		}
	}
	else
	{
		m_fs = sel;
	}
}

void GDT::AssumeGS(WORD sel)
{
	if (m_isActivated)
	{
		__asm
		{
			mov ax, sel
			mov gs, sel
		}
	}
	else
	{
		m_gs = sel;
	}
}

WORD GDT::AddCodeEntry(DWORD base, DWORD limit, BYTE ring)
{
	if (m_lastEntryId >= MAX_GDT_ENTRIES - 1)
	{
		return 0;
	}

	QWORD value =
	CreateEntry(base, limit, SEG_DESCTYPE(1) | SEG_PRES(1) | SEG_SAVL(0) | 
							 SEG_LONG(0)     | SEG_SIZE(1) | SEG_GRAN(1) | 
							 SEG_PRIV(ring)  | SEG_CODE_EXRD);

	m_lastEntryId++;

	m_entries[m_lastEntryId] = value;

	m_gdt.Limit = GDT_GET_LIMIT();

	return m_lastEntryId << 3;
}

WORD GDT::AddDataEntry(DWORD base, DWORD limit, bool isReadOnly, BYTE ring)
{
	if (m_lastEntryId >= MAX_GDT_ENTRIES - 1)
	{
		return 0;
	}

	QWORD value =
	CreateEntry(base, limit, SEG_DESCTYPE(1) | SEG_PRES(1) | SEG_SAVL(0) | 
							 SEG_LONG(0)     | SEG_SIZE(1) | SEG_GRAN(1) | 
							 SEG_PRIV(ring)  | (isReadOnly ? SEG_DATA_RD : SEG_DATA_RDWR));
	
	m_lastEntryId++;

	m_entries[m_lastEntryId] = value;
	
	m_gdt.Limit = GDT_GET_LIMIT();

	return m_lastEntryId << 3;
}